use std::net::TcpStream;
use std::io::{Read, Error, ErrorKind, Write};
use std::time::Duration;
use std::thread;
use std::sync::{Arc, Mutex};

const DATA_SLICE: usize = 1048576;
const TIME_SLICE: Option<Duration> = Some(Duration::from_millis(40));

pub struct CommandStream {
    tcp_stream: TcpStream
}

impl CommandStream {
    pub fn new(tcp_stream: TcpStream) -> CommandStream {
        CommandStream {
            tcp_stream
        }
    }

    pub fn read(&mut self) -> Result<String, std::io::Error> {
        let mut buf = [0u8];
        let mut bin = Vec::<u8>::new();
        loop {
            if self.tcp_stream.read(&mut buf)? < 1 {
                return Err(Error::new(ErrorKind::Interrupted,
                                      "Connection interrupted"));
            }
            if buf[0] != b'\n' {
                bin.push(buf[0]);
            } else {
                let command = match String::from_utf8(bin) {
                    Ok(c) => c,
                    Err(_) => {
                        return Err(Error::new(ErrorKind::InvalidData,
                                              "Command must be UTF-8 encoding"));
                    }
                };
                return Ok(command);
            }
        }
    }

    pub fn write(&mut self, command: String) -> Result<(), std::io::Error> {
        self.tcp_stream.write(command.as_bytes())?;
        self.tcp_stream.write(b"\n")?;
        Ok(())
    }

    pub fn set_read_timeout(&mut self, timeout: Option<Duration>) -> std::io::Result<()> {
        self.tcp_stream.set_read_timeout(timeout)
    }

    pub fn set_write_timeout(&mut self, timeout: Option<Duration>) -> std::io::Result<()> {
        self.tcp_stream.set_write_timeout(timeout)
    }

    pub fn try_clone_stream(&self) -> std::io::Result<TcpStream> {
        self.tcp_stream.try_clone()
    }
}

pub struct DataStream {
    tcp_stream: Arc<Mutex<TcpStream>>,
    head: Option<String>
}

impl DataStream {
    pub fn new(tcp_stream: TcpStream) -> DataStream {
        DataStream {
            tcp_stream: Arc::new(Mutex::new(tcp_stream)),
            head: None
        }
    }

    pub fn read_head(&mut self) -> Result<String, std::io::Error> {
        if let Some(head) = &self.head {
            return Ok(head.clone());
        }
        let mut buf = [0u8];
        let mut bin = Vec::<u8>::new();
        loop {
            let mut stream = self.tcp_stream.lock().unwrap();
            if stream.read(&mut buf)? < 1 {
                return Err(Error::new(ErrorKind::Interrupted,
                                      "Connection interrupted"));
            }
            if buf[0] != b'\n' {
                bin.push(buf[0]);
            } else {
                let head = match String::from_utf8(bin) {
                    Ok(c) => c,
                    Err(_) => {
                        return Err(Error::new(ErrorKind::InvalidData,
                                              "Command must be UTF-8 encoding"));
                    }
                };
                self.head = Some(head.clone());
                return Ok(head);
            }
        }
    }

    pub fn write_head(&mut self, head: String) -> Result<(), std::io::Error> {
        if let Some(_) = &self.head {
            return Err(Error::new(ErrorKind::InvalidData,
                                  "Cannot write head any more"));
        }
        self.head = Some(head.clone());
        let mut stream = self.tcp_stream.lock().unwrap();
        stream.write(head.as_bytes())?;
        stream.write(b"\n")?;
        Ok(())
    }

    pub fn pipe(&mut self, aim_stream: TcpStream) {
        let stop_flag = Arc::new(Mutex::new(false));
        let stop_flag_i = stop_flag.clone();

        let mutex = self.tcp_stream.clone();

        // Proxy Client Read Stream
        let my_stream = {
            let this_stream = mutex.lock().unwrap();
            this_stream.try_clone().unwrap()
        };

        // Proxy Client Write Stream
        let my_copy = my_stream.try_clone().unwrap();

        // Real Client Read Stream
        let aim_copy = aim_stream.try_clone().unwrap();

        // `aim_stream` is Real Client Write Stream

        // Proxy client pipe thread
        thread::spawn(move || {
            let stop_flag = stop_flag;
            let mut read_stream = my_stream;
            read_stream.set_read_timeout(TIME_SLICE).unwrap();
            let mut aim = aim_stream;
            let mut buffer = Box::new([0_u8; DATA_SLICE]);
            'outer: loop {
                let flag = {
                    let locked_flag = stop_flag.lock().unwrap();
                    *locked_flag
                };
                if flag { break 'outer; }
                loop {
                    match read_stream.read(buffer.as_mut()) {
                        Ok(size) => {
                            if size == 0 {
                                {
                                    let mut locked_flag = stop_flag.lock().unwrap();
                                    *locked_flag = true;
                                }
                                break 'outer;
                            }
                            match aim.write(&buffer.as_ref()[0..size]) {
                                Ok(wr_size) => {
                                    if wr_size < size {
                                        break 'outer;
                                    }
                                    continue;
                                }
                                Err(_) => {
                                    break 'outer;
                                }
                            };
                        }
                        Err(e) => {
                            match e.kind() {
                                ErrorKind::TimedOut | ErrorKind::WouldBlock => {
                                    continue 'outer;
                                }
                                _ => {
                                    break 'outer;
                                }
                            }
                        }
                    }
                }
            }
        });

        // Write thread
        thread::spawn(move || {
            let stop_flag = stop_flag_i;
            let mut write_stream = my_copy;
            let mut aim = aim_copy;
            aim.set_read_timeout(TIME_SLICE).unwrap();
            let mut buffer = Box::new([0_u8; DATA_SLICE]);
            'outer: loop {
                let flag = {
                    let locked_flag = stop_flag.lock().unwrap();
                    *locked_flag
                };
                if flag { break 'outer; }
                loop {
                    match aim.read(buffer.as_mut()) {
                        Ok(size) => {
                            if size == 0 {
                                {
                                    let mut locked_flag = stop_flag.lock().unwrap();
                                    *locked_flag = true;
                                }
                                break 'outer;
                            }
                            match write_stream.write(&buffer.as_ref()[0..size]) {
                                Ok(size) => {
                                    if size == 0 {
                                        break 'outer;
                                    }
                                    continue;
                                }
                                Err(_) => {
                                    break 'outer;
                                }
                            };
                        }
                        Err(e) => {
                            match e.kind() {
                                ErrorKind::TimedOut | ErrorKind::WouldBlock => {
                                    continue 'outer;
                                }
                                _ => {
                                    break 'outer;
                                }
                            }
                        }
                    }
                }
            }
        });
    }
}
